# @Time   : 2022/3/29 14:50
# @Author : Fioman
# @Phone  : 13149920693
# @Tips   : Talk is Cheap,Show me the code! ^_^^_^
"""
轮廓查找新算法,物料区域是白色情况.思路就是先找到上下的轮廓,然后找中间的额黑色图像.黑色图像查找之前要先将轮廓定义处理一下.找到最大的那个.
"""

# @Time   : 2022/3/29 9:52
# @Author : Fioman
# @Phone  : 13149920693
# @Tips   : Talk is Cheap,Show me the code! ^_^^_^


import time
import os
import cv2 as cv
import imutils
import numpy as np
from tools.config_params import ConfigParams
from vision_process.image_transfer import *

areaFilterUpDown = 500
areaFilterLeftRight = 150000
cp = ConfigParams()
colorBlue = (255, 0, 0)
colorGreen = (0, 255, 0)
colorRed = (0, 0, 255)
lineThickness = 5


def image_show(name, image, showState=True):
    """
    显示图片,用来调试的时候显示图片,按照一定的比例进行缩放
    :param name:显示的图片窗口名称
    :param image:显示的图片
    :param showState:是否显示
    :return:
    """
    if showState:
        sizeScale = 3
        newWidth = int(image.shape[1] / sizeScale)
        newHeight = int(image.shape[0] / sizeScale)
        showImage = cv.resize(image, (newWidth, newHeight), cv.INTER_AREA)
        cv.namedWindow("{}".format(name))
        cv.imshow("{}".format(name), showImage)
        cv.waitKey(0)


def get_size_scale_image(imageSrc, sizeScale):
    """
    高宽同等比例缩放图像,缩放比例
    :param imageSrc: 原图像
    :param sizeScale: 缩放比例
    :return:
    """
    h, w = imageSrc.shape[:2]
    newWidth, newHeight = int(w / sizeScale), int(h / sizeScale)
    newImage = cv.resize(imageSrc, (newWidth, newHeight), cv.INTER_AREA)
    return newImage


def get_center_by_contour(contour):
    """
    根据轮廓点获取这个轮廓的中心点
    :param contour:
    :return:
    """
    return int(cv.minAreaRect(contour)[0][0]), int(cv.minAreaRect(contour)[0][1])


def get_two_points_middle(p1, p2):
    """
    获取两个点的中点坐标
    :param p1: 点1
    :param p2: 点2
    :return:
    """
    assert len(p1) == 2
    assert len(p2) == 2
    middleX = (p1[0] + p2[0]) / 2
    middleY = (p1[1] + p2[1]) / 2
    return round(middleX, 2), round(middleY, 2)


def get_line_k_and_b(p1, p2):
    """
    湖区某一条直线的斜率和斜距
    :param p1: 点1的坐标
    :param p2: 点2的坐标
    :return:
    """
    if p1[0] == p2[0]:
        # 垂直于x轴,没有k值
        return None, p1[0]
    else:
        k = (p2[1] - p1[1]) / (p2[0] - p1[0])
        b = p1[1] - k * p1[0]
        return k, b


def get_two_lines_intersection(k1, b1, k2, b2):
    """
    获取平面上两条线的交点
    :param k1: 线1的k值
    :param b1: 线1的b值
    :param k2: 线2的k值
    :param b2: 线2的b值
    :return:
    """
    # 如果k相等,说明垂直无交点
    if k1 == k2:
        return None, None
    elif k1 is None and k2 is not None:
        return int(b1), int(k2 * b1 + b2)
    elif k2 is None and k1 is not None:
        return int(b2), int(k1 * b2 + b1)
    else:
        x = (b1 - b2) / (k1 - k2) * (-1)
        y = k1 * x + b1
    return int(x), int(y)


def get_rect_points_clockwise(rectPoints):
    """
    根据矩形的四个顶点,获取其顺时针方向的四个顶点
    :param rectPoints:
    :return:
    """
    rectPointsSorted = sorted(rectPoints, key=np.sum)
    leftTop = rectPointsSorted[0]
    rightBottom = rectPointsSorted[-1]

    if rectPointsSorted[1][0] > rectPointsSorted[2][0]:
        rightTop = rectPointsSorted[1]
        leftBottom = rectPointsSorted[2]
    else:
        rightTop = rectPointsSorted[2]
        leftBottom = rectPointsSorted[1]

    return tuple(leftTop), tuple(rightTop), tuple(rightBottom), tuple(leftBottom)


def get_linear_fit_k_and_b_vertical(pointsMat):
    """
    垂直方向获取拟合直线,因为垂直方向上,有角度可能很接近90度的情况,如果大于一定的值,就认为是None
    :param pointsMat:
    :return:
    """
    outPut = cv.fitLine(pointsMat, cv.DIST_L2, 0, 0.01, 0.01)
    if outPut[0] != 0:
        k = outPut[1] / outPut[0]
        k = k[0]
        if abs(k) > 1000:
            k = None
            b = np.mean(pointsMat, axis=0)[0]
        else:
            b = outPut[3] - k * outPut[2]
    else:
        k = None
        b = outPut[2]
    return float(k) if k is not None else None, float(b)


def find_total_board_info(image, width, height):
    cv.destroyAllWindows()
    image_show("BoardOriginal", image)
    # 将灯管之外的位置填黑
    image[1735:,:] = 0
    image[:70,:] = 0
    thresUsedVal = 180
    T, board = cv.threshold(image, thresUsedVal, 255, cv.THRESH_BINARY)
    image_show("BoardThresOriginal", board)

    kernel = cv.getStructuringElement(cv.MORPH_RECT, (5, 5))
    board = cv.morphologyEx(board, cv.MORPH_CLOSE, kernel, iterations=5)
    image_show("BoardOpenned", board)

    # 然后查找轮廓,找到轮廓的特征.按照轮廓的中心去归类,分为左右和上下.

    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)

    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt) > areaFilterUpDown]
    # 查找轮廓中心,按照轮廓的位置记录对应轮廓所属的可能的点位.
    centerLeftXMax = 756
    centerRightXmin = 1144
    centerTopPosMax = 100  # 竖着的轮廓按照第一边的位置,如果第一遍的位置小于100,就按照是上轮廓.如果大于100,就按照是下轮廓
    centerTopDownXMin = 874
    centerTopDownXMax = 950

    upContourList = []
    downContourList = []
    leftLine = None
    rightLine = None
    upPoint = None
    downPoint = None
    for cnt in largeContours:
        centerPoint = get_center_by_contour(cnt)
        centerX, centerY = centerPoint
        boxPoints = cv.boxPoints(cv.minAreaRect(cnt))
        rectPoints = get_rect_points_clockwise(boxPoints)
        leftTop, rightTop = rectPoints[:2]
        topPointY = get_two_points_middle(leftTop, rightTop)[1]
        if centerX < centerLeftXMax:
            continue
        elif centerX > centerRightXmin:
            continue
        elif topPointY < centerTopPosMax:
            if centerTopDownXMin < centerX < centerTopDownXMax:
                upContourList.append(cnt)
        else:
            if centerTopDownXMin < centerX < centerTopDownXMax:
                downContourList.append(cnt)

    # 先按照上下的轮廓去解析.如果上下都没有找到轮廓,就按照错误的去处理.如果找到一个轮廓,就按照推测去推另外一个轮廓,如果是全白的.找到的这个轮廓的面积很大,就报错.noBoard.
    # 急停没有检测到板子,识别失败,竖直灯管没有找到轮廓
    if len(upContourList) == 0 and len(downContourList) == 0:
        return "error", "竖直灯管没有找到任何轮廓", None, None, image

    if len(upContourList) >= 1 and len(downContourList) == 0:
        # 如果只找到了上轮廓,要看看这个轮廓的面积是多大,如果这个面积很大,占据了整个灯管,就认为是无板.
        upContour = sorted(upContourList, key=cv.contourArea, reverse=True)[0]
        contourArea = cv.contourArea(upContour)
        if contourArea > 110000:  # 如果竖直的灯管全部是量的,就认为是无板的情况
            return "noBoard", "竖直灯管没有任何遮挡", None,None, image
        else:
            # 这里是检测到上轮廓,然后下轮廓没找到,要推测下轮廓的位置.根据高度去推测
            upBox = cv.minAreaRect(upContour)
            upRectPoints = np.int0(cv.boxPoints(upBox))
            upClockwisePoint = get_rect_points_clockwise(upRectPoints)
            upLine = (upClockwisePoint[2], upClockwisePoint[3])
            downLine = (upLine[0][0], upLine[0][1] + height), (upLine[1][0], upLine[1][1] + height)
            imageColor = cv.cvtColor(image, cv.COLOR_GRAY2BGR)
            cv.line(imageColor, upLine[0], upLine[1], colorGreen, lineThickness)
            cv.line(imageColor, downLine[0], downLine[1], colorGreen, lineThickness)
            image_show("CalcResByUpContourAndHeight", imageColor)
            upPoint = get_two_points_middle(upLine[0], upLine[1])
            downPoint = get_two_points_middle(downLine[0], downLine[1])

    if len(upContourList) == 0 and len(downContourList) >= 1:
        downContour = sorted(downContourList, key=cv.contourArea, reverse=True)[0]
        contourArea = cv.contourArea(downContour)
        if contourArea > 110000:
            return "noBoard", "竖直灯管没有任何遮挡", None, None, image
        else:
            downBox = cv.minAreaRect(downContour)
            downRectPoints = np.int0(cv.boxPoints(downBox))
            downClockwisePoint = get_rect_points_clockwise(downRectPoints)
            downLine = (downClockwisePoint[0], downClockwisePoint[1])
            upLine = (downLine[0][0], downLine[0][1] - height), (downLine[1][0], downLine[1][1] - height)
            imageColor = cv.cvtColor(image, cv.COLOR_GRAY2BGR)
            cv.line(imageColor, upLine[0], upLine[1], colorGreen, lineThickness)
            cv.line(imageColor, downLine[0], downLine[1], colorGreen, lineThickness)
            image_show("CalcResByDownContourAndHeight", imageColor)
            upPoint = get_two_points_middle(upLine[0], upLine[1])
            downPoint = get_two_points_middle(downLine[0], downLine[1])

    if len(upContourList) >= 1 and len(downContourList) >= 1:
        upContour = sorted(upContourList, key=cv.contourArea, reverse=True)[0]
        downContour = sorted(downContourList, key=cv.contourArea, reverse=True)[0]
        upBox = cv.minAreaRect(upContour)
        upRectPoints = np.int0(cv.boxPoints(upBox))
        upClockwisePoint = get_rect_points_clockwise(upRectPoints)
        downBox = cv.minAreaRect(downContour)
        downRectPoints = np.int0(cv.boxPoints(downBox))
        downClockwisePoint = get_rect_points_clockwise(downRectPoints)
        upLine = upClockwisePoint[2], upClockwisePoint[3]
        downLine = downClockwisePoint[0], downClockwisePoint[1]
        imageColor = cv.cvtColor(image, cv.COLOR_GRAY2BGR)
        cv.line(imageColor, upLine[0], upLine[1], colorGreen, lineThickness)
        cv.line(imageColor, downLine[0], downLine[1], colorGreen, lineThickness)
        image_show("CalcResByUpContourAndDownContour", imageColor)
        upPoint = get_two_points_middle(upLine[0], upLine[1])
        downPoint = get_two_points_middle(downLine[0], downLine[1])

    # 将upPoint 和 downPoint附近的区域填充成黑色,再找一遍轮廓,填充的像素是200,差不多100个mm
    boardFillBlack = 200
    if upPoint[1] + boardFillBlack < board.shape[0]:
        board[:int(upPoint[1] + boardFillBlack), :] = 255

    board[1553:, :] = 255
    T, board = cv.threshold(board, 150, 255, cv.THRESH_BINARY_INV)
    image_show("BoardFindedForLeftRight", board)

    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)

    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt) > areaFilterLeftRight]
    # 上下蹲额轮廓找到以后,就是看左右的轮廓.如果左右的轮廓都没有找到,就报错noBoard
    contourOkList = []
    if len(largeContours) == 0:
        return "error", "没有找到任何轮廓", None, None, image
    else:
        print("找到的轮廓面积: {}".format([cv.contourArea(contour) for contour in largeContours]))
        # 剔除掉无效轮廓,就是那种在边界的轮廓.
        for contour in largeContours:
            centerX, centerY = get_center_by_contour(contour)
            if 100 < centerX < 1820:
                contourOkList.append(contour)
    if len(contourOkList) == 0:
        return "error", "没有找到任何轮廓", None, None, image

    # 轮廓找到以后,找竖直的直线.
    scharrX = cv.Scharr(board,cv.CV_64F,1,0)
    scharrY = cv.Scharr(board,cv.CV_64F,0,1)
    scharrX = cv.convertScaleAbs(scharrX)
    scharrY = cv.convertScaleAbs(scharrY)
    scharrXY = cv.addWeighted(scharrX,0.2,scharrY,0.2,0)
    board = scharrXY.copy()
    image_show("ScharrXY",board)
    leftXMin,leftXMax = 1,2


    boardContour = sorted(contourOkList, key=cv.contourArea, reverse=True)[0]
    boardBox = cv.minAreaRect(boardContour)
    boardRectPoints = get_rect_points_clockwise(np.int0(cv.boxPoints(boardBox)))

    leftTop, rightTop, rightBottom, leftBottom = boardRectPoints  # 上下左右四个点.
    leftLineMin = min(leftTop[0], leftBottom[0])
    rightLineMax = max(rightTop[0], rightTop[0])
    # 如果左边的线和右边的线太靠近边缘,就认为找到的轮廓不对
    if leftLineMin < 5 and rightLineMax + 5 < board.shape[1]:
        # 通过右边的线推测左边的线.
        rightLine = rightTop, rightBottom
        rightK, rightB = get_line_k_and_b(rightLine[0], rightLine[1])
        if rightK is None:
            # 如果是垂线,直接延长
            leftLine = (rightLine[0][0] - width, rightLine[0][1]), (rightLine[1][0] - width, rightLine[1][1])
        else:
            # 先求出来垂线的斜率
            verK = -1 / rightK
            xOffset = np.cos(trans_k_to_rad(verK)) * width
            leftLine = (rightLine[0][0] - xOffset, rightLine[0][1]), (rightLine[1][0] - xOffset, rightLine[1][1])
    elif leftLineMin > 5 and rightLineMax + 5 > board.shape[1]:
        # 通过左边的额线条推测右边的线条
        leftLine = leftTop, leftBottom
        leftK, leftB = get_line_k_and_b(leftLine[0], leftLine[1])
        if leftK is None:
            rightLine = (leftLine[0][0] + width, leftLine[0][1]), (leftLine[1][0] + width, leftLine[1][1])
        else:
            verK = -1 / leftK
            xOffset = np.cos(trans_k_to_rad(verK)) * width
            rightLine = (leftLine[0][0] + xOffset, leftLine[0][1]), (leftLine[1][0] + xOffset, leftLine[1][1])
    elif leftLineMin > 5 and rightLineMax + 5 < board.shape[1]:
        leftLine = leftTop, leftBottom
        rightLine = rightTop, rightBottom
    else:
        return "error", "没有找到任何轮廓", None, None, image

    # 现在是四个轮廓都找到了,对应的点位也找到了.
    leftLineK, leftLineB = get_line_k_and_b(leftLine[0], leftLine[1])
    rightLineK, rightLineB = get_line_k_and_b(rightLine[0], rightLine[1])

    # 获取四个交点.然后根据交点画线.
    interSection1 = get_ver_line_intersection(leftLineK, leftLineB, upPoint)
    interSection2 = get_ver_line_intersection(rightLineK, rightLineB, upPoint)
    interSection3 = get_ver_line_intersection(rightLineK, rightLineB, downPoint)
    interSection4 = get_ver_line_intersection(leftLineK, leftLineB, downPoint)

    # 画出轮廓
    imageColor = cv.cvtColor(image, cv.COLOR_GRAY2BGR)
    imageColorContour = imageColor.copy()

    cv.line(imageColorContour, interSection1, interSection2, (0, 255, 0), 3)
    cv.line(imageColorContour, interSection2, interSection3, (0, 255, 0), 3)
    cv.line(imageColorContour, interSection3, interSection4, (0, 255, 0,), 3)
    cv.line(imageColorContour, interSection4, interSection1, (0, 255, 0), 3)
    image_show("ContourFind", imageColorContour)

    # 根据左右的斜率和斜距,获取角平分线,求物料角度
    meanLineK, meanLineB = get_mean_angle(leftLineK, leftLineB, rightLineK, rightLineB)
    angle = get_angle_by_k(meanLineK)

    # 获取抓取位置偏移
    yOffset = int(cp.catchPosOffset * cp.mm2pix)

    fixK, fixB = get_line_k_and_b(interSection3, interSection4)

    if meanLineK is None:
        fixMiddle = get_two_points_middle(interSection3, interSection4)
        catchX = fixMiddle[0]
        catchY = fixMiddle[1] - yOffset
    else:
        fixMiddle = get_two_points_middle(interSection3, interSection4)
        # catchX = fixMiddle[0] - abs(yOffset * np.cos(trans_k_to_rad(meanLineK)))
        # catchY = meanLineK * catchX + meanLineB
        catchY = fixMiddle[1] - abs(yOffset * np.sin(trans_k_to_rad(meanLineK)))
        catchX = (catchY - meanLineB) / meanLineK
    if angle < 0:
        angle = round(angle + 90, 2)
    else:
        angle = round(angle - 90, 2)

    catchPos = (int(catchX), int(catchY))
    # 和固定边的交点
    interSection5 = get_two_lines_intersection(fixK, fixB, meanLineK, meanLineB)
    # (x,y)坐标颠倒
    # interSection5 = (interSection5[1],interSection5[0])

    cv.line(imageColorContour, interSection5, catchPos, colorRed, lineThickness)
    cv.circle(imageColorContour, catchPos, 5, colorRed, -1)
    cv.putText(imageColorContour, "angle:{:.2f}".format(angle), (catchPos[0] - 100, catchPos[1] - 200),
               cv.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 4)
    cv.putText(imageColorContour, "catchPos:{}".format(catchPos), (catchPos[0] - 100, catchPos[1] - 30), cv.FONT_HERSHEY_SIMPLEX,
               2, (255, 0, 0), 4)
    print("Angle: {:.2f},catchPos:{}".format(angle, catchPos))
    image_show("TotalBoard", imageColorContour)

    # 根据图像抓取位置,获取实际的坐标偏移
    catchPos = get_module_catch_pos(catchPos, (image.shape[1] // 2, image.shape[0] // 2))
    if -100 < catchPos[0] < 970 and 0 < catchPos[1] < 240:  # 底层限位,有可能会算错
        return "ok", "识别成功!", catchPos, angle, imageColorContour
    else:
        return "ok", "识别的坐标超过极限位置", catchPos, angle, imageColorContour


def get_module_catch_pos(catchPos, imageCenter):
    # 抓取的图像位置
    catchXpix, catchYPix = catchPos
    # 标定的图像中心对应的模组坐标
    imageCenterXPix, imageCenterYPix = imageCenter
    imageOffsetX = (catchXpix - imageCenterXPix) / cp.mm2pix  # 转换成mm
    imageOffsetY = (catchYPix - imageCenterYPix) / cp.mm2pix  # 转换成mm

    # 实际的抓板位置
    catchPosX = cp.centerModuleX - imageOffsetY
    catchPosY = cp.centerModuleY - imageOffsetX
    return round(catchPosX, 2), round(catchPosY, 2)


def trans_k_to_rad(k):
    """
    将斜率转换为弧度
    :param k: 斜率k值
    :return:
    """
    return np.arctan(k)


def get_angle_by_k(k):
    """
    根据斜率k值,获取对应的角度值
    :param k: 线条的斜率k值
    :return:
    """
    if k == 0:
        return 0
    if k is None:
        return 90
    return np.arctan(k) * 180 / np.pi


def get_mean_angle(k1, b1, k2, b2):
    """
    获取角平分线
    :param k1:
    :param b1:
    :param k2:
    :param b2:
    :return:
    """

    intersection = get_two_lines_intersection(k1, b1, k2, b2)
    # 如果本身就是平行线没有交点.这样中线的k就是upK 或者downK,都一样,如果确定b的值呢.
    if intersection[0] is None:
        middleLineK = k1
        x1 = b1
        x2 = b2
        if middleLineK is not None:
            # 先找到中线上的一个必过点
            middleLineP = ((x1+x2) / 2, 0)
            middleLineB = middleLineP[1] - middleLineK * middleLineP[0]
        else:
            middleLineK = None
            middleLineB = (x1 + x2) / 2
    else:
        middleLineK = np.tan((np.arctan(k1) + np.arctan(k2)) / 2)  # 角平分线公式
        middleLineB = intersection[1] - middleLineK * intersection[0]

    return middleLineK, middleLineB


def get_ver_line_intersection(k1, b1, point):
    """
    获取一个点和一条直线的垂直交点
    :param k1:
    :param b1:
    :param point:
    :return:
    """
    if k1 is None:
        intersection = (int(b1), int(point[1]))
    elif k1 == 0:
        intersectionX = int(point[0])
        intersectionY = int(b1)
        intersection = (intersectionX, intersectionY)
    else:
        verK = -1 / k1
        verB = point[1] - verK * point[0]
        intersection = get_two_lines_intersection(k1, b1, verK, verB)
    return intersection


def distance(p1, p2):
    """
    计算平面上两个点之前的距离,精确到小数点后两位
    :param p1: 平面上的点1
    :param p2: 平面上的点2
    :return:
    """
    assert len(p1) == 2
    assert len(p2) == 2
    return round(np.sqrt((abs(p2[0] - p1[0]) ** 2 + abs(p2[1] - p1[1]) ** 2)), 2)


if __name__ == '__main__':
    filePath = r"D:\kehui_plate\original_no_scale\fail\ok"
    keepDir = filePath
    if not os.path.isdir(keepDir):
        keepDir, _ = os.path.split(filePath)

    keepOkPath = os.path.join(keepDir, "ok")
    keepOkResPath = os.path.join(keepDir, "res")
    keepFailPath = os.path.join(keepDir, "fail")
    if not os.path.exists(keepOkPath):
        os.makedirs(keepOkPath)
    if not os.path.exists(keepOkResPath):
        os.makedirs(keepOkResPath)
    if not os.path.exists(keepFailPath):
        os.makedirs(keepFailPath)

    width = 1174
    height = 1550

    fileNames = []  # 保存要去识别的所有的图片
    if os.path.isfile(filePath):
        filePath, fileName = os.path.split(filePath)
        fileNames.append(fileName)
    else:
        for root, dirs, files in os.walk(filePath):
            if root != filePath:
                continue
            for file in files:
                if file[-4:] != ".bmp":
                    continue
                fileNames.append(file)

    for index, fileName in enumerate(fileNames):
        print("图像名称:{} ".format(fileName))
        filePathReal = os.path.join(filePath, fileName)
        imageTest = cv.imread(filePathReal, cv.IMREAD_GRAYSCALE)
        # imageUsed = get_image_used(imageTest)
        # cv.imwrite(os.path.join(keepOkPath,fileName),imageUsed)
        # continue
        print(imageTest.shape)
        calcStart = time.time()
        res, info, center, angle, imageRes = find_total_board_info(imageTest, width, height)
        if res == "ok":
            print("算法耗时: {}".format(time.time() - calcStart))
            os.remove(filePathReal)
            imageRes = get_size_scale_image(imageRes, sizeScale=4)
            cv.imwrite(os.path.join(keepOkResPath, fileName), imageRes)
            cv.imwrite(os.path.join(keepOkPath, fileName), imageTest)
            print("识别成功")
            continue
        else:
            os.remove(filePathReal)
            cv.imwrite(os.path.join(keepFailPath, fileName), imageTest)
            print("算法耗时: {}".format(time.time() - calcStart))
            print("识别失败")
